package cn.cloudtogo.common.jwt;

import cn.cloudtogo.common.GendbException;
import com.auth0.jwt.JWT;
import com.auth0.jwt.JWTCreator;
import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.exceptions.JWTDecodeException;
import com.auth0.jwt.exceptions.JWTVerificationException;
import com.auth0.jwt.interfaces.DecodedJWT;
import com.auth0.jwt.interfaces.JWTVerifier;
import com.auth0.jwt.interfaces.Verification;
import lombok.extern.slf4j.Slf4j;

import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author yaolianhua789@gmail.com
 **/
@Slf4j
public class GendbJwt implements Jwt {

    private final List<String> officialDefinition = List.of("iss", "sub", "exp", "nbf", "iat", "aud", "jti");

    @Override
    public String signHS256(JwtClaims jwtClaims, String secret) {
        HeaderClaims headerClaims = jwtClaims.headerClaims();
        PayLoadClaims payLoadClaims = jwtClaims.payLoadClaims();
        Algorithm algorithm = Algorithm.HMAC256(secret);
        JWTCreator.Builder builder = JWT.create();
        if (headerClaims != null) {
            builder.withKeyId(headerClaims.getKeyId());
        }
        if (payLoadClaims != null) {
            builder.withIssuer(payLoadClaims.getIssuer())
                    .withAudience(payLoadClaims.getAudience().toArray(new String[]{}))
                    .withJWTId(payLoadClaims.getId())
                    .withSubject(payLoadClaims.getSubject())
                    .withExpiresAt(payLoadClaims.getExpiresAt())
                    .withNotBefore(payLoadClaims.getNotBefore())
                    .withIssuedAt(payLoadClaims.getIssuedAt());

            Map<String, String> data = payLoadClaims.getData();
            if (data != null) {
                data.forEach(builder::withClaim);
            }
        }

        return builder.sign(algorithm);
    }

    @Override
    public boolean verifyHS256(String token, String secret, Map<String, String> claims) {
        try {
            Algorithm algorithm = Algorithm.HMAC256(secret);
            Verification verification = JWT.require(algorithm);
            if (claims != null) {
                claims.forEach(verification::withClaim);
            }

            JWTVerifier verifier = verification.build();
            verifier.verify(token);
        } catch (JWTVerificationException exception) {
            //Invalid signature/claims
            return false;
        }
        return true;
    }

    @Override
    public JwtClaims decodeToken(String token) {
        try {
            DecodedJWT jwt = JWT.decode(token);
            HeaderClaims headerClaims = HeaderClaims.builder()
                    .type(jwt.getType())
                    .keyId(jwt.getKeyId())
                    .algorithm(jwt.getAlgorithm())
                    .build();
            Map<String, String> data = new HashMap<>(16);
            jwt.getClaims().entrySet().stream()
                    .filter(e -> !officialDefinition.contains(e.getKey()))
                    .forEach(entry -> data.put(entry.getKey(), entry.getValue().asString()));
            PayLoadClaims payLoadClaims = PayLoadClaims.builder()
                    .expiresAt(jwt.getExpiresAt())
                    .audience(jwt.getAudience())
                    .id(jwt.getId())
                    .issuedAt(jwt.getIssuedAt())
                    .issuer(jwt.getIssuer())
                    .notBefore(jwt.getNotBefore())
                    .subject(jwt.getSubject())
                    .data(data)
                    .build();

            return new JwtClaims() {
                @Override
                public HeaderClaims headerClaims() {
                    return headerClaims;
                }

                @Override
                public PayLoadClaims payLoadClaims() {
                    return payLoadClaims;
                }
            };
        } catch (JWTDecodeException e) {
            throw new GendbException("invalid token");
        }
    }

    @Override
    public String signRS256(JwtClaims jwtClaims, RSAPublicKey publicKey, RSAPrivateKey privateKey) {
        Algorithm algorithm = Algorithm.RSA256(publicKey, privateKey);
        HeaderClaims headerClaims = jwtClaims.headerClaims();
        PayLoadClaims payLoadClaims = jwtClaims.payLoadClaims();

        JWTCreator.Builder builder = JWT.create();
        if (headerClaims != null) {
            builder.withKeyId(headerClaims.getKeyId());
        }
        if (payLoadClaims != null) {
            builder.withIssuer(payLoadClaims.getIssuer())
                    .withAudience(payLoadClaims.getAudience().toArray(new String[]{}))
                    .withJWTId(payLoadClaims.getId())
                    .withSubject(payLoadClaims.getSubject())
                    .withExpiresAt(payLoadClaims.getExpiresAt())
                    .withNotBefore(payLoadClaims.getNotBefore())
                    .withIssuedAt(payLoadClaims.getIssuedAt());

            Map<String, String> data = payLoadClaims.getData();
            if (data != null) {
                data.forEach(builder::withClaim);
            }
        }

        return builder.sign(algorithm);
    }

    @Override
    public boolean verifyRS256(String token, RSAPublicKey publicKey, RSAPrivateKey privateKey, Map<String, String> claims) {
        try {
            Algorithm algorithm = Algorithm.RSA256(publicKey, privateKey);
            Verification verification = JWT.require(algorithm);
            if (claims != null) {
                claims.forEach(verification::withClaim);
            }

            JWTVerifier verifier = verification.build();
            verifier.verify(token);
        } catch (JWTVerificationException exception) {
            //Invalid signature/claims
            return false;
        }
        return true;
    }
}
